import torch
from torch import nn


class AlexNet(nn.Module):
    def __init__(self):
        super(AlexNet, self).__init__()
        self.features = nn.Sequential(
            # input 3*227*227  (227 - 11) / 4 + 1 = 55  output 96*55*55
            nn.Conv2d(3, 96, kernel_size=11, stride=4),
            nn.ReLU(inplace=True),
            # input 96*55*55  (55 - 3) / 2 + 1 = 27  output 96*27*27
            nn.MaxPool2d(kernel_size=3, stride=2),
            # input 96*27*27  27 - 5 + 2 * 2 + 1 = 27  output 256*27*27
            nn.Conv2d(96, 256, kernel_size=5, padding=2),
            nn.ReLU(inplace=True),
            # input256*27*27  (27 - 3) / 2 + 1 = 13  output 256*13*13
            nn.MaxPool2d(kernel_size=3, stride=2),
            # input 256*13*13  13 - 3 + 1 * 2 + 1 = 13  output 384*13*13
            nn.Conv2d(256, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 384, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            # input 256*13*13  (13 - 3) / 2 + 1 = 6  output 256*6*6
            nn.MaxPool2d(kernel_size=3, stride=2)
        )

        self.classifier = nn.Sequential(
            # 减少过拟合
            nn.Dropout(p=0.5),
            # 全连接
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(p=0.5),
            nn.Linear(4096, 4096),
            nn.ReLU(inplace=True),
            nn.Linear(4096, 1000)
        )

    def forward(self, x):
        # x:[batch_size, 256, 6, 6]
        x = self.features(x)
        # x:[batch_size, 9216]
        x = x.view(x.shape[0], -1)
        x = self.classifier(x)
        return x
